# -*- coding = utf-8 -*-
# @Time : 2022/3/9 18:39
# @Author : GHHHHHHHHH
# @File : funksvdInference.py
# @Software : PyCharm
import torch
from train import TrainFunkSVD
import numpy as np

# demo
class FunksvdInference:
    def __init__(self, mode='train', dataset=None, hidden=None, epoch=50, LR=0.01):
        self.__model = None
        self.__mode = mode
        self.__dataset = dataset
        self.__hidden = hidden
        self.__epoch = epoch
        self.__LR = LR

    def __getModel(self):
        if self.__mode == 'train':
            if self.__dataset is None:
                raise ValueError("Dataset can't be None if you want to train the model!")
            trainer = TrainFunkSVD(self.__dataset.shape[0], self.__dataset.shape[1], self.__dataset, hidden=self.__hidden
                                   , epoch=self.__epoch, LR=self.__LR)
            trainer.train()
        elif self.__mode == 'load':
            pass
        else:
            raise ValueError("Parameter 'mode' can just be 'train' or 'load'!")
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = torch.load('model.pth').to(device)

    def __getSVD(self):
        self.__getModel()
        return self.model().data

    def get_list(self, user_id, length):
        arr = self.__getSVD().cpu().numpy().tolist()
        arr = [i[user_id] for i in arr]
        return_list = []
        while return_list.__len__() < length:
            return_list.append(arr.index(max(arr)))
            arr[arr.index(max(arr))] = -100000000
        return return_list

if __name__ == '__main__':
    # --------------------Demo----------------------------
    # Generate random matrix as training dataset and input it as numpy matrix, abscissa user and ordinate item
    arr1 = np.random.rand(50, 20)
    # dataset，which means user-item array；
    # hidden: number of hidden layers，if hidden == None, it will looking for the most excellent parameters，but need longer time;
    # mode can just be 'train' or 'load'
    tool = FunksvdInference(dataset=arr1, hidden=100, mode='train', epoch=200)
    # return the recommendation list generated by parameter user_id and length
    print(tool.get_list(user_id=0, length=10))


